/*!
*  Copyright (c) 2015 by Joshua Zhang
* \file detector.hpp
* \brief ssd detection module impl
*/

#include "zupply.hpp"
#include "genderpredictor.hpp"
#include <iostream>
#include <sstream>
#include <streambuf>
#include <cassert>
#include <cstdlib>
#include <ctime>

using namespace zz;

namespace FaceInfo {
	GenderAge::GenderAge(const std::vector<std::string>& model_zoo, int epoch, int width,
		int height, float mean_r, float mean_g, float mean_b,
		int device_type, int device_id)
	{
		auto logger = log::get_logger("default");
		if (epoch < 0 || epoch > 9999)
		{
			logger->error("Invalid epoch number: ") << epoch;
			exit(-1);
		}

		if (model_zoo.empty()||model_zoo.size()>3)
		{
			logger->error("Invalid model: ");
			exit(-1);
		}

		for (int i = 0; i < model_zoo.size(); i++)
		{
			std::string model_file = model_zoo[i] + "-" + fmt::int_to_zero_pad_str(epoch, 4) + ".params";
			if (!os::is_file(model_file)) {
				std::cerr << "Model file: " << model_file << " does not exist" << std::endl;
				exit(-1);
			}
			std::string json_file = model_zoo[i] + "-symbol.json";
			if (!os::is_file(json_file)) {
				std::cerr << "JSON file: " << model_file << " does not exist" << std::endl;
				exit(-1);
			}
			if (width < 1 || height < 1) {
				std::cerr << "Invalid width or height: " << width << "," << height << std::endl;
				exit(-1);
			}
			width_ = width;
			height_ = height;
			const char *input_name = "data";
			const char *input_keys[1];
			input_keys[0] = input_name;
			const mx_uint input_shape_indptr[] = { 0, 4 };
			const mx_uint input_shape_data[] = { 1, 3, static_cast<mx_uint>(width_), static_cast<mx_uint>(height_) };
			mean_r_ = mean_r;
			mean_g_ = mean_g;
			mean_b_ = mean_b;

			// load model
			std::ifstream param_file(model_file, std::ios::binary | std::ios::ate);
			if (!param_file.is_open()) {
				std::cerr << "Unable to open model file: " << model_file << std::endl;
				exit(-1);
			}
			std::streamsize size = param_file.tellg();
			param_file.seekg(0, std::ios::beg);
			buffer_.resize(size);

			std::ifstream json_handle(json_file, std::ios::ate);
			std::string json_buffer;
			json_buffer.reserve(json_handle.tellg());
			json_handle.seekg(0, std::ios::beg);
			json_buffer.assign((std::istreambuf_iterator<char>(json_handle)), std::istreambuf_iterator<char>());
			if (json_buffer.size() < 1) {
				std::cerr << "invalid json file: " << json_file << std::endl;
				exit(-1);
			}

			if (param_file.read(buffer_.data(), size))
			{

				if (i == 0) {
					MXPredCreate(json_buffer.c_str(), buffer_.data(), static_cast<int>(size),
						device_type, device_id, 1, input_keys, input_shape_indptr, input_shape_data,
						&predictor_);
				}
				else if (i == 1) {
					MXPredCreate(json_buffer.c_str(), buffer_.data(), static_cast<int>(size),
						device_type, device_id, 1, input_keys, input_shape_indptr, input_shape_data,
						&predictor_female);
				}
				else if (i == 2) {
					MXPredCreate(json_buffer.c_str(), buffer_.data(), static_cast<int>(size),
						device_type, device_id, 1, input_keys, input_shape_indptr, input_shape_data,
						&predictor_male);
				}
			}
			else {
				std::cerr << "Unable to read model file: " << model_file << std::endl;
				exit(-1);
			}
		}

	}

	cv::Mat GenderAge::prepare_face(cv::Mat image, const cv::Rect& rect)const
	{
		int x1 = rect.x;
		int y1 = rect.y;
		int x2 = rect.x + rect.width;
		int y2 = rect.y + rect.height;
		int face_width = x2 - x1;
		int face_height = y2 - y1;
		int new_x1 = x1 - int(face_width*0.2);
		int new_y1 = y1 - int(face_height*0.5);
		int new_x2 = x2 + int(face_width*0.2);
		int new_y2 = y2 + int(face_height*0.15);
		if (new_x1 < 0)	new_x1 = 0;
		if (new_y1 < 0) new_y1 = 0;
		if (new_x2 > image.cols) new_x2 = image.cols;
		if (new_y2 > image.rows) new_y2 = image.rows;

		return image(cv::Range(new_y1, new_y2), cv::Range(new_x1, new_x2)).clone();
	}

	int GenderAge::predictface(cv::Mat image, const cv::Rect& rect) {
		
		cv::Mat face = prepare_face(image,rect);
		std::vector<float>&& dets = predict(face,0).first;
		int result = std::distance(dets.begin(), std::max_element(dets.begin(), dets.end()));
		return result;
	}

	std::pair <std::vector<float>, std::vector<float> > GenderAge::predict(const cv::Mat& image, int flag)
	{
		auto logger = log::get_logger("default");
		if (image.empty()) {
			std::cerr << "Image is empty" << std::endl;
			exit(-1);
		}
		if (image.channels() != 3) {
			std::cerr << "RGB image required" << std::endl;
			exit(-1);
		}
		// prepare data
		cv::Mat resize_img;
		int width_ = 224;
		int height_ = 224;
		cv::resize(image, resize_img, cv::Size(width_, height_));
		int size = resize_img.channels() * width_ * height_;
		cv::Mat image_float;
		resize_img.convertTo(image_float, CV_32FC3);
		cv::Mat bgr[3];
		cv::split(image_float, bgr);
		bgr[0] = bgr[0] - mean_b_;
		bgr[1] = bgr[1] - mean_g_;
		bgr[2] = bgr[2] - mean_r_;
		cv::Mat stack = bgr[2];
		stack.push_back(bgr[1]);
		stack.push_back(bgr[0]);
		// use model to forward
		mx_uint *shape = NULL;
		mx_uint shape_len = 0;
		std::vector<float> outputs, outputs_age;
		if (flag == 0) {
			MXPredSetInput(predictor_, "data", (float*)stack.data, static_cast<mx_uint>(size));
			time::Timer timer;
			MXPredForward(predictor_);
			MXPredGetOutputShape(predictor_, 0, &shape, &shape_len);
			mx_uint tt_size = 1;
			for (mx_uint i = 0; i < shape_len; ++i) {
				tt_size *= shape[i];
			}
			outputs.resize(tt_size);
			MXPredGetOutput(predictor_, 0, outputs.data(), tt_size);
			//logger->info("Gender forward elapsed time: ") << timer.to_string();
		}
		else if (flag == 1)
		{
			MXPredSetInput(predictor_female, "data", (float*)stack.data, static_cast<mx_uint>(size));
			time::Timer timer;
			MXPredForward(predictor_female);
			MXPredGetOutputShape(predictor_female, 0, &shape, &shape_len);
			mx_uint tt_size = 1;
			for (mx_uint i = 0; i < shape_len; ++i) {
				tt_size *= shape[i];
			}
			outputs_age.resize(tt_size);
			MXPredGetOutput(predictor_female, 0, outputs_age.data(), tt_size);
			//logger->info("Age forward elapsed time: ") << timer.to_string();
		}

		else if (flag == 2)
		{
			MXPredSetInput(predictor_male, "data", (float*)stack.data, static_cast<mx_uint>(size));
			time::Timer timer;
			MXPredForward(predictor_male);
			MXPredGetOutputShape(predictor_male, 0, &shape, &shape_len);
			mx_uint tt_size = 1;
			for (mx_uint i = 0; i < shape_len; ++i) {
				tt_size *= shape[i];
			}
			outputs_age.resize(tt_size);
			MXPredGetOutput(predictor_male, 0, outputs_age.data(), tt_size);
			//logger->info("Age forward elapsed time: ") << timer.to_string();
		}
		return std::make_pair(outputs, outputs_age);
	}

	std::pair <int, std::vector<float>> GenderAge::predictface_with_age(cv::Mat image, const cv::Rect& rect) {
		cv::Mat face = prepare_face(image,rect);
		std::vector<float> dets = predict(face, 0).first;
		int result = std::distance(dets.begin(), std::max_element(dets.begin(), dets.end()));
		int gender_flag = result + 1;
		std::vector<float>&& dets_age = predict(face, gender_flag).second;
		//int age_result = std::distance(dets_age.begin(), std::max_element(dets_age.begin(), dets_age.end()));
		return std::make_pair(result, dets_age);
	}
}  // namespace det
